#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2018 crane <crane@crane-pc>

from functools import reduce

from math_tools import float_equal
from linear_tools import linear_combination, add_vector,  is_zero_vector, scale_vector
from dot_cross_product import dot_product
import copy


def size(mat):
    return len(mat), len(mat[0])

def diagonal(mat):
    for i in range(len(mat)):
        yield mat[i][i]

def iter_matrix_cols(matrix):
    cols = len(matrix[0])

    for col in range(cols):
        the_col_vector = [i[col] for i in matrix]
        yield the_col_vector

def iter_matrix_rows(matrix):
    '''matrix本身就可以用来遍历row
       for row in matrix
    '''
    return matrix

def col_of(matrix, col_idx):
    return [row[col_idx] for row in iter_matrix_rows(matrix)]
    # for row in iter_matrix_rows(matrix):
    #     yield row[col_idx]

def linear_combination_col(matrix, col_vector):
    return linear_combination(iter_matrix_cols(matrix), col_vector)

def linear_combination_row(row_vector, matrix):
    return linear_combination(iter_matrix_rows(matrix), row_vector)
    # return linear_combination(matrix, row_vector)

def scale_matrix(matrix, co):
    results = []
    for row in iter_matrix_rows(matrix):
        new_row = scale_vector(row, co)
        results.append(new_row)

    return results

# ===================  =====================
def resort_mat_row(mat, new_idxs):
    # print('new idxs %s' % new_idxs)
    new_mat = [mat[row] for row in new_idxs]
    return new_mat

# def rotate_matrix(matrix):
def transpose_matrix(matrix):
    ''' 矩阵转置m[i,j] = mt[j,i] '''
    results = []
    for col in iter_matrix_cols(matrix):
        results.append(list(col))
    return results

def identity(nth):
    '''生成一个nth identity 矩阵'''
    one_row = [0] * nth
    identity = [list(one_row) for i in range(nth)]

    for i in range(nth):
        identity[i][i] = 1

    return identity

def trace(mat):
    rows, cols = size(mat)
    assert rows == cols, (rows, cols)
    return sum( diagonal(mat) )

def multiply(*mats):
    return reduce(multiply_1, mats)

def multiply_1(m1, m2):
    '''normal method
        从结果element角度.
    '''
    row = len(m1)
    col = len(m2[0])
    one_row = [0] * col
    result = [list(one_row) for r in range(row)]

    for r in range(row):
        for c in range(col):
            # row =
            the_col = [ i[c] for i in m2]
            result[r][c] = dot_product(m1[r], the_col)

    return result

def multiply_2(m1, m2):
    '''
        matrix * col_vector (linear combination)
    '''
    results = []
    for col in iter_matrix_cols(m2):
        result_col = linear_combination_col(m1, col)
        results.append(list(result_col))

    # results是列向量组合, 变为行向量组合矩阵(rotate)
    return transpose_matrix(results)

def multiply_3(m1, m2):
    '''
        row_vector * matrix (linear combination)
    '''
    results = []

    for row in iter_matrix_rows(m1):
        result_row = linear_combination_row(row, m2)
        results.append(list(result_row))

    return results

def multiply_4(m1, m2):
    '''
        sum(A's row * B' col):
        直观理解: 先计算所有的x未知数, 再计算所有的y未知数, 再计算z...
        然后结果矩阵相加
    '''
    mid_sum = None

    iter_cols = iter_matrix_cols(m1)    # 第一列是所有的x 的系数(第二列是y系数)
    iter_rows = iter_matrix_rows(m2)    # 第一行时所有的x变量值(第二行是y变量值)

    def sum(c, r):
        nonlocal mid_sum
        c = list(c)
        r = list(r)
        r = [r]
        c = transpose_matrix([c])
        s = multiply_1(c, r)
        if not mid_sum:
            mid_sum = s
        else:
            mid_sum = matrix_sum(mid_sum, s)

    list(map(sum, list(iter_cols), list(iter_rows)))

    return mid_sum


# def common_power(base, nth, func, identity_ele):
#     ''' identity_ele 幺元'''

#     assert type(nth) is int
#     remain = nth % 2
#     half = nth // 2
#     return

def matrix_power(mat, nth):
    ''' return mat^nth '''
    if nth == 0:
        return identity(len(mat))
    if nth == 1:
        return mat

    remain = nth % 2
    half = nth // 2

    half_power = matrix_power(mat, half)
    even_power = multiply_1(half_power, half_power)

    if remain:
        return multiply_1(even_power, mat)
    else:
        return even_power

def matrix_sum(m1, m2):
    ''' 两个矩阵相加 '''
    results = []
    m1_rows = iter_matrix_rows(m1)
    m2_rows = iter_matrix_rows(m2)
    def row_add(r1, r2):
        row_s = list(add_vector(r1, r2))
        results.append(row_s)

    list(map(row_add, m1_rows, m2_rows))
    return results

def matrix_append_col(m1, m2):
    results = []
    def append(r1, r2):
        results.append(r1 + r2)

    list(map(append, m1, m2))
    return results

def matrix_append_row(m1, m2, inplace=False):
    results = m1
    if not inplace:
        results = copy.deepcopy(m1)

    for row in m2:
        results.append(row)

    return results

def matrix_slice_row(mat, row_idxs):
    ''' [row_idx:] '''
    return [mat[row] for row in row_idxs]

def matrix_slice_col(mat, col_idx):
    ''' [row_idx:] '''

def kernel(mat):
    ''' Null(A) space representation '''

def matrix_equal(a, b):
    rows, cols = size(a)
    if size(a) != size(b):
        return False

    for row in range(rows):
        for col in range(cols):
            value1 = a[row][col]
            value2 = b[row][col]
            # assert value1 == value2,  [value1, value2, row, col]
            if not float_equal(value1, value2):
                return False

    return True

def is_zero_matrix(mat):
    for row in mat:
        if not is_zero_vector(row):
            return False

    return True

def is_symmetric(mat):
    ''' 1: 矩阵沿对角线(1,1 2,2 3,3 4,4)对称.
        2: 矩阵的m(i,j) == m(j,i).
        3: 矩阵的转置transpose矩阵, 还是矩阵本身.
    '''
    rows = len(mat)
    cols = len(mat[0])

    if rows != cols:
        return False

    for i in range(rows):
        for j in range(i):
            v1 = mat[i][j]
            v2 = mat[j][i]
            if not float_equal(v1, v2):
                return False

    return True

def is_identity(mat):
    ''' 1: 矩阵沿对角线(1,1 2,2 3,3 4,4)对称.
        2: 矩阵的m(i,j) == m(j,i).
        3: 矩阵的转置transpose矩阵, 还是矩阵本身.
    '''
    rows = len(mat)
    cols = len(mat[0])

    if rows != cols:
        return False

    for r in range(rows):
        for c in range(cols):
            v = mat[r][c]
            if r != c:
                if not float_equal(v, 0):
                    return False
            else:
                if not float_equal(v, 1):
                    return False

    return True


def copy_matrix(mat):
    ''' deep copy a matrix '''
    return copy.deepcopy(mat)

def show_matrix(mat, msg=''):
    print('=================== mat: %s =====================' % msg)
    for r in mat:
        print(r)
    print()

# =================== test/application  =====================

def show(mat, msg=''):
    print('=================== mat: %s =====================' % msg)
    for r in mat:
        print(r)

def test():
    test_multiply()
    test_power()
    test_transpose()
    test_trace()
    test_matrix_sum()
    test_matrix_append()

def test_power():
    fib_matrix = [
        [0, 1],     # 0*a1 + 1*a2 = new_a1 (0, 1)
        [1, 1],     # 1*a1 + 1*a2 = new_a2 (1, 1)
    ]
    r = matrix_power(fib_matrix, 3)
    show(r, 'matrix_power')

def test_transpose():
    ''' 乘的结果的转置, 等于逆序转置的乘
        A*B = M
        T(M) = T(B) * T(A)
        可以从: 列向量线性组合(右乘) 转换为 行向量线性组合(左乘)角度建立直觉
    '''
    A = [
        [3, 5, 8],
        [4, 6, 1],
        [2, 3, 5],
    ]
    B = [
        [5, 5, 8],
        [8, 6, 1],
        [9, 3, 15],
    ]

    AB = multiply_1(A, B)
    ABT = transpose_matrix(AB)

    BT = transpose_matrix(B)
    AT = transpose_matrix(A)
    BT_AT = multiply_1(BT, AT)
    show(AB, 'AB')
    show(ABT, 'ABT')
    show(BT_AT, 'BT_AT == ABT')

def test_matrix_sum():
    A = [
        [3, 5, 8],
        [4, 6, 1],
        [2, 3, 5],
    ]
    B = [
        [5, 5, 8],
        [8, 6, 1],
        [9, 3, 15],
    ]
    show(matrix_sum(A, B), 'MATRIX_SUM_A_B')

def test_matrix_append():
    A = [
        [3, 5, 8],
        [4, 6, 1],
        [2, 3, 5],
    ]
    B = [
        [5, 5, 8],
        [8, 6, 1],
        [9, 3, 15],
    ]
    show(matrix_append_col(A, B), 'matrix_append_col')

def test_trace():
    A = [
        [3, 5, 8],
        [4, 6, 1],
        [2, 3, 5],
    ]

    v = trace(A)
    print('\ntrace of A %s' % v)


def test_multiply():
    m1 = [[1, 2, ], [ 5, 6], [ 8, 9]]
    m2 = [[1, 2, 3], [3, 5, 6]]

    # r1 = multiply_1(m1, m1)
    r1 = multiply_1(m1, m2)
    show(r1, 'multiply_1')

    # r2 = linear_combination_col(m1, [2, 4])
    # print(list(r2))
    rc = linear_combination_col(m1, [1, 2, 3])
    print(list(rc))

    r2 = multiply_2(m1, m2)
    show(r2, 'multiply_2')

    r3 = multiply_3(m1, m2)
    show(r3, 'multiply_3')

    r4 = multiply_4(m1, m2)
    show(r4, 'multiply_4')


    m1 = [
        [1, 2],
        [5, 6],
    ]
    rm = multiply(m1, m1, m1)
    show(rm, 'multiply 3 times')


def main():
    print("start main")
    test()

if __name__ == "__main__":
    main()
